"""
Pipeline example demonstrating advanced routing features
"""
import os
from typing import Dict, List
from sonetto.agents.user import UserProxy
from sonetto.agents.chat import ChatAgent
from sonetto.agents.tool import ToolAgent
from sonetto.core.scheduler import Scheduler
from sonetto.core.toolset import McpToolset
from sonetto.core.message import Message, MessageType
from sonetto.core.routing import CommonConditions, CommonBroadcastConditions

def get_weather(city: str) -> str:
    """模拟天气查询工具"""
    return f"{city}的天气是晴朗，温度25℃"

def main():
    # 获取API配置
    api_key = os.getenv("API_KEY")
    if not api_key:
        raise ValueError("请设置API_KEY环境变量")
    
    base_url = os.getenv("API_BASE_URL")

    # 创建调度器
    scheduler = Scheduler()

    # 创建用户代理
    user = UserProxy(agent_id="user")
    scheduler.register_agent(user)

    # 创建解析代理
    parser = ChatAgent(
        agent_id="parser",
        api_key=api_key,
        base_url=base_url,
        system_prompt="你是一个解析器，负责分析用户的意图并添加相应的标签。"
    )
    scheduler.register_agent(parser)

    # 创建工具集
    toolset = McpToolset()

    # 注册工具
    @toolset.register_tool(
        name="get_weather",
        description="查询指定城市的天气情况"
    )
    def weather_tool(city: str) -> str:
        return get_weather(city)

    # 创建工具代理
    tool_agent = ToolAgent(
        agent_id="tool_assistant",
        api_key=api_key,
        base_url=base_url,
        toolset=toolset
    )
    scheduler.register_agent(tool_agent)

    # 创建回复生成代理
    response_agent = ChatAgent(
        agent_id="response_generator",
        api_key=api_key,
        base_url=base_url,
        system_prompt="你是一个回复生成器，负责生成友好的回复。"
    )
    scheduler.register_agent(response_agent)

    # 创建监控代理（接收所有消息的副本）
    monitor = ChatAgent(
        agent_id="monitor",
        api_key=api_key,
        base_url=base_url,
        system_prompt="你是一个监控器，负责记录和分析所有消息。"
    )
    scheduler.register_agent(monitor)

    # 设置处理流程顺序
    scheduler.set_agent_order([
        "user",
        "parser",
        "tool_assistant",
        "response_generator",
        "user"
    ])

    # 添加条件路由
    scheduler.add_conditional_route(
        source="parser",
        condition=CommonConditions.has_metadata_tag("requires_tool"),
        target="tool_assistant",
        description="需要工具处理"
    )

    scheduler.add_conditional_route(
        source="parser",
        condition=CommonConditions.has_metadata_tag("direct_response"),
        target="response_generator",
        description="直接生成回复"
    )

    # 添加广播规则
    scheduler.add_broadcast_pattern(
        source="user",
        targets=["monitor"],
        description="监控用户输入"
    )

    scheduler.add_broadcast_pattern(
        source="tool_assistant",
        targets=["monitor"],
        condition=CommonConditions.message_type_is("TOOL_RESULT"),
        description="监控工具执行结果"
    )

    # 输出路由图
    print("当前路由配置：")
    print("```mermaid")
    print(scheduler.visualize_routes())
    print("```")

    print("\n=== Sonetto 流水线处理示例 ===")
    print("可用命令：")
    print("- 输入任何消息开始对话")
    print("- 输入 'exit' 退出")
    print("=============================")

    while True:
        try:
            user_input = input("\n你: ").strip()
            if user_input.lower() == 'exit':
                break

            # 创建初始消息
            message = Message(
                type=MessageType.TEXT,
                content=user_input,
                sender=user.agent_id
            )

            # 开始消息处理流程
            response = scheduler.step(message)
            while response:
                response = scheduler.step(response)

        except KeyboardInterrupt:
            break
        except Exception as e:
            print(f"发生错误: {str(e)}")

    # 导出路由配置（可以保存到文件中供后续使用）
    config = scheduler.export_config()
    print("\n路由配置：")
    print(config.to_json(indent=2))

if __name__ == "__main__":
    main()
